package db

import (
	"encoding/json"

	"github.com/aquasecurity/trivy-db/pkg/types"

	bolt "go.etcd.io/bbolt"
	"golang.org/x/xerrors"
)

const (
	vulnerabilityDetailBucket = "vulnerability-detail"
)

func (dbc Config) PutVulnerabilityDetail(tx *bolt.Tx, cveID, source string, vuln types.VulnerabilityDetail) error {
	root, err := tx.CreateBucketIfNotExists([]byte(vulnerabilityDetailBucket))
	if err != nil {
		return err
	}
	return dbc.put(root, cveID, source, vuln)
}

// TODO: Until we can dependency inject, we need to monkey patch sadly
var (
	dbcForEachFunc = Config{}.forEach
)

func (dbc Config) GetVulnerabilityDetail(cveID string) (map[string]types.VulnerabilityDetail, error) {
	values, err := dbcForEachFunc(vulnerabilityDetailBucket, cveID)
	if err != nil {
		return nil, xerrors.Errorf("error in NVD get: %w", err)
	}
	if len(values) == 0 {
		return nil, nil
	}

	vulns := map[string]types.VulnerabilityDetail{}
	for source, value := range values {
		var vuln types.VulnerabilityDetail
		if err = json.Unmarshal(value, &vuln); err != nil {
			return nil, xerrors.Errorf("failed to unmarshal Vulnerability JSON: %w", err)
		}
		vulns[source] = vuln
	}
	return vulns, nil
}

func (dbc Config) DeleteVulnerabilityDetailBucket() error {
	return dbc.deleteBucket(vulnerabilityDetailBucket)
}
